from datetime import datetime

from langchain_core.messages import AnyMessage
from langchain_core.runnables import RunnableConfig
from langchain_core.tools import tool
from langgraph.prebuilt import create_react_agent
from langgraph.prebuilt.chat_agent_executor import AgentState

from agent.my_llm import llm
from agent.my_state import CustomState
from agent.tools.agent_state_tools import get_username, greet_user


@tool("get_weather")
def get_weather(city: str, date: str) -> str:
    """
    获取某个城市的天气
    参数：
        city-城市
        date-日期
    """
    print("call get weather")
    return f'日期：{date}，城市：{city}的气温为35℃，天气晴转多云。'


@tool("get_current_date")
def get_current_date() -> str:
    """
    获取今天日期
    """
    return datetime.today().strftime("%Y-%m-%d")


@tool("get_user_info")
def get_user_info(config: RunnableConfig) -> dict:
    """
    获取用户所有信息，包括：姓名，性别，年龄等
    """
    print('call get user info')
    username = config['configurable'].get('username', 'zhangsan')
    print(f'调用工具传入的用户名是：{username}')
    # 根据用户名，从数据源中查询用户信息
    return {'username': username, 'name': 'Alm', 'age': 18, 'gender': '男'}


# 系统提示词模板函数
# 由用户传入内容，组成一个动态的系统提示词
def prompt(state: AgentState, config: RunnableConfig) -> list[AnyMessage]:
    username = config['configurable'].get('username', 'zhangsan')
    print(f'username: {username}')
    system_message = f'你是一个智能AI助手，你的任务是回答用户的问题，必要时请调用工具获得答案。如果包含问题涉及【今天】或者【明天】这些词汇，先用get_current_date这个工具查出今天的日期，再使用get_weather这个工具查询天气。当前用户的名字是：{username}'
    return [{"role": "system", "content": system_message}] + state['messages']


# 注意：这里agent必须名字为graph，否则会找不到
graph = create_react_agent(
    model=llm,
    # tools=[get_current_date, get_weather, get_user_info],
    tools=[get_current_date, get_weather, get_username, greet_user],
    # tools=[],
    prompt='你是一个智能AI助手，你的任务是回答用户的问题，必要时请调用工具获得答案。如果包含问题涉及【今天】或者【明天】这些词汇，先用get_current_date这个工具查出今天的日期，再使用get_weather这个工具查询天气。生成祝福语请调用相应的tool',
    # prompt=prompt,
    state_schema=CustomState,
)
